#!/usr/bin/env python
# -*- coding: utf-8 -*-

from datetime import datetime
from typing import List, Dict, Any, Optional, Literal
from pydantic import BaseModel, Field, validator


# 项目模型
class ProjectBase(BaseModel):
    """项目基础模型"""
    name: str = Field(..., description="项目名称")
    description: Optional[str] = Field(None, description="项目描述")


class ProjectCreate(ProjectBase):
    """项目创建模型"""
    pass


class ProjectUpdate(BaseModel):
    """项目更新模型"""
    name: Optional[str] = Field(None, description="项目名称")
    description: Optional[str] = Field(None, description="项目描述")


class ProjectResponse(ProjectBase):
    """项目响应模型"""
    id: str = Field(..., description="项目ID")
    created_at: datetime = Field(..., description="创建时间")
    updated_at: datetime = Field(..., description="更新时间")
    
    class Config:
        orm_mode = True
        from_attributes = True


# 训练数据模型
class TrainingDataBase(BaseModel):
    """训练数据基础模型"""
    type: Literal["ddl", "question_sql", "description"] = Field(..., description="数据类型")
    content: str = Field(..., description="内容")
    question: Optional[str] = Field(None, description="问题")
    sql: Optional[str] = Field(None, description="SQL查询")
    
    @validator("question")
    def validate_question(cls, v, values):
        """验证问题字段"""
        if values.get("type") == "question_sql" and not v:
            raise ValueError("问题字段对于question_sql类型是必需的")
        return v
    
    @validator("sql")
    def validate_sql(cls, v, values):
        """验证SQL字段"""
        if values.get("type") == "question_sql" and not v:
            raise ValueError("SQL字段对于question_sql类型是必需的")
        return v


class TrainingDataCreate(TrainingDataBase):
    """训练数据创建模型"""
    pass


class TrainingDataResponse(TrainingDataBase):
    """训练数据响应模型"""
    id: int = Field(..., description="训练数据ID")
    project_id: str = Field(..., description="项目ID")
    created_at: datetime = Field(..., description="创建时间")
    
    class Config:
        orm_mode = True
        from_attributes = True


class TrainingDataFilter(BaseModel):
    """训练数据过滤模型"""
    type: Optional[Literal["ddl", "question_sql", "description"]] = Field(None, description="数据类型")
    limit: int = Field(10, description="限制数量")
    offset: int = Field(0, description="偏移量")


# 文本到SQL模型
class Text2SQLRequest(BaseModel):
    """文本到SQL请求模型"""
    question: str = Field(..., description="用户问题")
    db_type: str = Field("mysql", description="数据库类型")
    system_prompt: Optional[str] = Field(None, description="自定义系统提示词")


class Text2SQLResponse(BaseModel):
    """文本到SQL响应模型"""
    sql: str = Field(..., description="生成的SQL查询")


# SQL纠错模型
class ErrorAttempt(BaseModel):
    """错误尝试模型"""
    sql: str = Field(..., description="错误的SQL查询")
    error_message: str = Field(..., description="错误信息")


class SQLCorrectionRequest(BaseModel):
    """SQL纠错请求模型"""
    question: str = Field(..., description="原始问题")
    error_attempts: List[ErrorAttempt] = Field(..., description="错误尝试列表")
    db_type: str = Field("mysql", description="数据库类型")
    system_prompt: Optional[str] = Field(None, description="自定义系统提示词")


class SQLCorrectionResponse(BaseModel):
    """SQL纠错响应模型"""
    sql: str = Field(..., description="纠正后的SQL查询")


# 查询历史模型
class QueryHistoryBase(BaseModel):
    """查询历史基础模型"""
    type: Literal["generation", "correction"] = Field(..., description="历史类型")
    question: str = Field(..., description="问题")
    generated_sql: str = Field(..., description="生成的SQL")
    error_sql: Optional[str] = Field(None, description="错误的SQL")
    error_message: Optional[str] = Field(None, description="错误信息")
    db_type: str = Field("mysql", description="数据库类型")


class QueryHistoryCreate(QueryHistoryBase):
    """查询历史创建模型"""
    pass


class QueryHistoryResponse(QueryHistoryBase):
    """查询历史响应模型"""
    id: int = Field(..., description="历史ID")
    project_id: str = Field(..., description="项目ID")
    created_at: datetime = Field(..., description="创建时间")
    
    class Config:
        orm_mode = True
        from_attributes = True


class QueryHistoryFilter(BaseModel):
    """查询历史过滤模型"""
    type: Optional[Literal["generation", "correction"]] = Field(None, description="历史类型")
    start_date: Optional[datetime] = Field(None, description="开始日期")
    end_date: Optional[datetime] = Field(None, description="结束日期")
    db_type: Optional[str] = Field(None, description="数据库类型")
    limit: int = Field(10, description="限制数量")
    offset: int = Field(0, description="偏移量")


# 通用响应模型
class MessageResponse(BaseModel):
    """消息响应模型"""
    message: str = Field(..., description="消息内容")


class ErrorResponse(BaseModel):
    """错误响应模型"""
    error: Dict[str, Any] = Field(..., description="错误信息")


class PaginatedResponse(BaseModel):
    """分页响应模型"""
    items: List[Any] = Field(..., description="项目列表")
    total: int = Field(..., description="总数")
    limit: int = Field(..., description="限制数量")
    offset: int = Field(..., description="偏移量")